# 前景主体分割算法 - 操作界面
from GUI.UI_Window import Ui_MainWindow
from Info_Pane import Info_Window
from dataset.dataset_image import get_loader
from solver import Solver

from PyQt5.QtCore import *
from PyQt5.QtWidgets import *
from PyQt5.QtGui import *
from torch.cuda import is_available
from shutil import copyfile
from cv2 import cvtColor, COLOR_BGR2RGB,imdecode
from numpy import fromfile,uint8
import argparse
import sys
import qdarkstyle

# 设置根目录路径
import os
root = os.path.abspath(os.path.dirname(__file__))
root = root.replace("\\", "/")


# 检测界面窗口
class Window(Ui_MainWindow, QMainWindow):
    # 初始化
    def __init__(self):
        super(Ui_MainWindow, self).__init__()
        self.setAttribute(Qt.WA_StyledBackground, True)
        self.setupUi(self)
        # 设置深色主题
        self.setStyleSheet(qdarkstyle.load_stylesheet_pyqt5())
        # 设置窗口图标
        self.setWindowIcon(QIcon(root + "/GUI/Images/icon.ico"))

        # 保存当前选择的图片路径
        self.image_path = None
        # 保存结果图片路径
        self.mask_path = root + "/results/mask.png"
        self.result_path = root + "/results/result.png"
        # 标记图片检测状态
        self.is_segmented = False

        # 设置参数选项
        self.config = self.setup_config()
        # 设置推理网络模型
        self.solver = Solver(self.config)

        # 按钮事件绑定
        self.Btn_Upload.clicked.connect(self.Upload_Image)
        self.Btn_Remove.clicked.connect(self.Remove_Image)
        self.Btn_Segment.clicked.connect(self.Segment_Image)
        self.Btn_Save.clicked.connect(self.Save_Image)
        # 菜单栏事件绑定
        self.Action_Upload.triggered.connect(self.Upload_Image)
        self.Action_Remove.triggered.connect(self.Remove_Image)
        self.Action_Segment.triggered.connect(self.Segment_Image)
        self.Action_Save.triggered.connect(self.Save_Image)
        self.Action_Help.triggered.connect(self.Show_Help)
        self.Action_Us.triggered.connect(self.About_Us)
        self.Action_Exit.triggered.connect(self.Exit_System)


    # 设置参数选项
    def setup_config(self):
        parser = argparse.ArgumentParser()
        if is_available():
            parser.add_argument('--cuda', type = bool, default = True)
        else:
            parser.add_argument('--cuda', type = bool, default = False)
        parser.add_argument('--arch', type = str, default = 'resnet')
        parser.add_argument('--input_size', type = int, default = 384)
        parser.add_argument('--model', type = str, default = root + '/final.pth')
        parser.add_argument('--test_on_gpus', type = list, default = [0])
        config = parser.parse_args()
        return config


    # 将图片显示在界面区域中
    def Show_Image(self, image, Label):
        image = cvtColor(image, COLOR_BGR2RGB)
        height, width, bytesPerComponent = image.shape
        bytesPerLine = bytesPerComponent * width
        q_image = QImage(image.data, width, height, bytesPerLine, QImage.Format_RGB888).scaled(Label.width(), Label.height())
        Label.setPixmap(QPixmap.fromImage(q_image))


    # 移除图片 - Btn_Remove 按钮响应
    def Remove_Image(self):
        # 若当前没有上传图片, 则不响应
        if self.image_path is None:
            QMessageBox.warning(self, "警告", "当前没有上传图片！", QMessageBox.Yes)
            return

        # 清除所有显示的图片
        self.Label_Figure.setText(" ")
        self.Label_Mask.setText(" ")
        self.Label_Result.setText(" ")
        # 重置标记变量
        self.image_path = None
        self.is_segmented = False


    # 上传图片 - Btn_Upload 按钮响应
    def Upload_Image(self):
        # 若已上传图片, 则要求先移除当前图片
        if self.image_path is not None:
            QMessageBox.warning(self, "警告", "请先移除当前图片！", QMessageBox.Yes)
            return

        # 打开文件选择对话框
        image_selected, _ = QFileDialog.getOpenFileName(self, "选择上传图片", "", "*.jpg;;*.png;;All Files(*)")

        # "" 表示用户取消选择
        if image_selected != "":
            #img = imread(image_selected)
            img = imdecode(fromfile(image_selected, dtype=uint8), -1)
            self.image_path = image_selected
            self.Show_Image(img, self.Label_Figure)


    # 主体分割 - Btn_Segment 按钮响应
    def Segment_Image(self):
        # 若未上传图片, 则要求先上传图片
        if self.image_path is None:
            QMessageBox.warning(self, "警告", "当前没有上传图片！", QMessageBox.Yes)
            return
        # 若已完成检测, 则不响应
        if self.is_segmented:
            QMessageBox.warning(self, "警告", "当前图片已完成主体分割！", QMessageBox.Yes)
            return

        # 设置读取图片的 dataloader
        test_loader = get_loader(self.image_path, self.config, mode = 'test')
        self.solver.setup_dataloader(test_loader)
        # 获取 mask 结果图片和 result 主体分割图片
        self.solver.generate_mask(self.mask_path)
        self.solver.generate_result(self.image_path, self.mask_path, self.result_path)

        # 读取并显示结果
        #mask = imread(self.mask_path)
        #result = imread(self.result_path)
        mask = imdecode(fromfile(self.mask_path, dtype=uint8), -1)
        result = imdecode(fromfile(self.result_path, dtype=uint8), -1)
        self.Show_Image(mask, self.Label_Mask)
        self.Show_Image(result, self.Label_Result)
        # 标记检测状态
        self.is_segmented = True


    # 保存图片 - Btn_Save 按钮响应
    def Save_Image(self):
        # 若未上传图片, 则要求先上传图片
        if self.image_path is None:
            QMessageBox.warning(self, "警告", "当前没有上传图片！", QMessageBox.Yes)
            return
        # 若图片未检测, 则要求先完成检测分割
        if not self.is_segmented:
            QMessageBox.warning(self, "警告", "请先完成主体分割！", QMessageBox.Yes)
            return

        # 打开路径选择对话框
        save_dir = QFileDialog.getExistingDirectory(self, "选择保存路径", "")
        # "" 表示用户取消选择
        if save_dir != "":
            # 复制结果图片至指定路径
            copyfile(self.mask_path, save_dir + "/mask.png")
            copyfile(self.result_path, save_dir + "/result.png")
            # 提示保存完成
            QMessageBox.warning(self, "提示", "结果已保存至指定路径！", QMessageBox.Yes)


    # 显示操作指南
    def Show_Help(self):
        info_window.hide()
        info_window.Help()
        info_window.show()

    # 显示团队信息
    def About_Us(self):
        info_window.hide()
        info_window.Us()
        info_window.show()

    # 安全退出系统
    def Exit_System(self):
        window.close()
        info_window.close()


# 主函数
if __name__ == '__main__':
    # 加载界面
    app = QApplication(sys.argv)
    info_window = Info_Window()
    window = Window()
    # 显示主界面
    info_window.hide()
    window.show()
    sys.exit(app.exec())